import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.optimize import minimize
import time
import logging

# 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger()

# 设置matplotlib使用支持中文的字体
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']  # 使用支持更广的字体
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

# 无人机参数
mass = 1.0  # kg
g = 9.81  # m/s^2
dt = 0.1  # 时间步长
T_horizon = 1.0  # 预测时域长度
N = int(T_horizon / dt)  # 预测步数

# 无人机几何参数
rotor_radius = 0.1  # m
body_radius = 0.2  # m
arm_length = 0.3  # m

# 空气动力学参数
drag_coeff = 0.1  # 阻力系数

# 状态约束 - 放宽约束
max_velocity = 8.0  # m/s (增加最大速度)
max_tilt = 45.0 * np.pi / 180.0  # 最大倾斜角度 (rad) (增加最大倾斜角)
min_thrust = 0.3 * mass * g  # 最小推力 (N) (降低最小推力)
max_thrust = 2.5 * mass * g  # 最大推力 (N) (增加最大推力)

# 目标点 - 更近的目标
target_pos = np.array([5.0, 5.0, 3.0])

# 障碍物 (位置和半径) - 减少障碍物数量
obstacles = [
    {"pos": np.array([2.0, 2.0, 1.5]), "radius": 0.8},
    {"pos": np.array([3.5, 4.0, 2.0]), "radius": 1.0}
]

# 简化版无人机动力学模型
def drone_dynamics(state, u):
    """
    简化版无人机动力学模型
    state: [x, y, z, vx, vy, vz, roll, pitch, yaw]
    u: [thrust, roll_torque, pitch_torque, yaw_torque]
    """
    x, y, z, vx, vy, vz, roll, pitch, yaw = state
    thrust, roll_torque, pitch_torque, yaw_torque = u
    
    # 简化的旋转矩阵
    R = np.array([
        [np.cos(yaw)*np.cos(pitch), 
         np.cos(yaw)*np.sin(pitch)*np.sin(roll) - np.sin(yaw)*np.cos(roll),
         np.cos(yaw)*np.sin(pitch)*np.cos(roll) + np.sin(yaw)*np.sin(roll)],
        [np.sin(yaw)*np.cos(pitch),
         np.sin(yaw)*np.sin(pitch)*np.sin(roll) + np.cos(yaw)*np.cos(roll),
         np.sin(yaw)*np.sin(pitch)*np.cos(roll) - np.cos(yaw)*np.sin(roll)],
        [-np.sin(pitch),
         np.cos(pitch)*np.sin(roll),
         np.cos(pitch)*np.cos(roll)]
    ])
    
    # 重力加速度
    gravity = np.array([0, 0, -g])
    
    # 推力加速度 (在机体Z轴方向)
    thrust_acc = R.dot(np.array([0, 0, thrust])) / mass
    
    # 空气阻力 (与速度平方成正比，方向相反)
    air_vel = np.array([vx, vy, vz])
    air_speed = np.linalg.norm(air_vel)
    if air_speed > 0.01:
        drag_force = -0.5 * drag_coeff * air_speed * air_vel
        drag_acc = drag_force / mass
    else:
        drag_acc = np.zeros(3)
    
    # 总加速度
    acceleration = gravity + thrust_acc + drag_acc
    
    # 状态导数
    dx = vx
    dy = vy
    dz = vz
    dvx, dvy, dvz = acceleration
    droll = roll_torque  # 简化为直接控制角速度
    dpitch = pitch_torque
    dyaw = yaw_torque
    
    return np.array([dx, dy, dz, dvx, dvy, dvz, droll, dpitch, dyaw])

# 离散化动力学模型 (欧拉法)
def discrete_dynamics(state, u, dt):
    return state + dt * drone_dynamics(state, u)

# 障碍物成本函数
def obstacle_cost(position, obstacles, safe_distance=1.5):
    """
    计算障碍物成本
    """
    cost = 0.0
    for obs in obstacles:
        dist = np.linalg.norm(position - obs["pos"])
        if dist < safe_distance + obs["radius"]:
            # 使用更平滑的惩罚函数
            penalty = 1.0 / (dist - obs["radius"] + 0.1)  # 避免除以零
            cost += 50.0 * penalty
    return cost

# 成本函数
def cost_function(u_flat, current_state, target_pos, obstacles):
    """
    计算NMPC成本函数，包含障碍物避障
    """
    # 重塑控制输入序列
    u_seq = u_flat.reshape((N, 4))
    
    # 初始化成本和状态
    cost = 0.0
    state = current_state.copy()
    
    # 预测未来N步
    for i in range(N):
        # 应用控制输入
        state = discrete_dynamics(state, u_seq[i], dt)
        
        # 位置误差成本
        pos_error = state[0:3] - target_pos
        cost += 5.0 * np.dot(pos_error, pos_error)  # 降低位置误差权重
        
        # 速度成本 (鼓励稳定)
        cost += 0.5 * np.dot(state[3:6], state[3:6])  # 降低速度权重
        
        # 姿态成本 (鼓励水平飞行)
        cost += 0.1 * (state[6]**2 + state[7]**2)  # 降低姿态权重
        
        # 控制输入成本
        cost += 0.05 * np.dot(u_seq[i], u_seq[i])  # 降低控制输入权重
        
        # 障碍物避障成本
        cost += obstacle_cost(state[0:3], obstacles)
    
    # 终端成本
    cost += 50.0 * np.dot(pos_error, pos_error)  # 降低终端成本权重
    
    return cost

# 约束函数 - 简化约束
def constraint_function(u_flat, current_state):
    """
    定义约束函数
    """
    # 重塑控制输入序列
    u_seq = u_flat.reshape((N, 4))
    
    # 初始化约束
    constraints = []
    state = current_state.copy()
    
    # 只检查第一步的约束，减少计算复杂度
    state = discrete_dynamics(state, u_seq[0], dt)
    
    # 速度约束
    velocity = np.linalg.norm(state[3:6])
    constraints.append(max_velocity - velocity)  # 速度约束: v <= max_velocity
    
    # 姿态约束
    roll, pitch = state[6:8]
    tilt = np.sqrt(roll**2 + pitch**2)
    constraints.append(max_tilt - tilt)  # 倾斜角约束: tilt <= max_tilt
    
    # 推力约束
    thrust = u_seq[0, 0]
    constraints.append(thrust - min_thrust)  # 推力约束: thrust >= min_thrust
    constraints.append(max_thrust - thrust)  # 推力约束: thrust <= max_thrust
    
    return np.array(constraints)

# NMPC控制器
def nmpc_controller(current_state, target_pos, obstacles):
    """
    NMPC控制器
    """
    # 初始猜测 (基于当前状态和目标方向)
    # 计算到目标的方向向量
    target_dir = target_pos - current_state[0:3]
    target_dist = np.linalg.norm(target_dir)
    if target_dist > 0.1:
        target_dir = target_dir / target_dist
    
    # 基于目标方向生成初始猜测
    base_thrust = mass * g
    u0 = np.tile([base_thrust, 0, 0, 0], N)
    
    # 添加一些随机性以避免陷入局部最优
    u0 += 0.1 * np.random.randn(len(u0))
    
    # 定义约束
    cons = {'type': 'ineq', 'fun': constraint_function, 'args': (current_state,)}
    
    # 优化问题 - 使用更宽松的优化参数
    result = minimize(cost_function, u0, args=(current_state, target_pos, obstacles),
                     constraints=cons, method='SLSQP', 
                     options={'maxiter': 30, 'disp': False, 'ftol': 1e-2, 'eps': 1e-2})
    
    # 计算预测轨迹
    predicted_trajectory = []
    if result.success:
        # 提取最优控制序列
        u_opt = result.x.reshape((N, 4))
        
        # 计算预测轨迹
        state_pred = current_state.copy()
        for i in range(N):
            state_pred = discrete_dynamics(state_pred, u_opt[i], dt)
            predicted_trajectory.append(state_pred[0:3].copy())
        
        return u_opt[0], np.array(predicted_trajectory)
    else:
        logger.warning("Optimization failed: %s", result.message)
        # 尝试使用更简单的控制策略
        # 计算到目标的方向
        target_dir = target_pos - current_state[0:3]
        target_dist = np.linalg.norm(target_dir)
        
        if target_dist > 0.1:
            target_dir = target_dir / target_dist
            
            # 简单的PD控制
            kp = 0.5
            kd = 0.1
            
            # 位置误差
            pos_error = target_pos - current_state[0:3]
            
            # 速度误差
            vel_error = -current_state[3:6]
            
            # 计算期望加速度
            desired_acc = kp * pos_error + kd * vel_error + np.array([0, 0, g])
            
            # 计算所需推力
            thrust = mass * np.linalg.norm(desired_acc)
            thrust = np.clip(thrust, min_thrust, max_thrust)
            
            # 计算姿态
            if thrust > 0.1:
                acc_dir = desired_acc / np.linalg.norm(desired_acc)
                pitch = np.arcsin(-acc_dir[0])
                roll = np.arcsin(acc_dir[1] / np.cos(pitch))
                
                # 限制姿态
                roll = np.clip(roll, -max_tilt, max_tilt)
                pitch = np.clip(pitch, -max_tilt, max_tilt)
                
                # 计算角速度控制
                roll_rate = 0.5 * (roll - current_state[6])
                pitch_rate = 0.5 * (pitch - current_state[7])
                yaw_rate = 0.0  # 保持当前偏航角
                
                # 为简单控制策略生成预测轨迹
                state_pred = current_state.copy()
                u_simple = np.array([thrust, roll_rate, pitch_rate, yaw_rate])
                for i in range(N):
                    state_pred = discrete_dynamics(state_pred, u_simple, dt)
                    predicted_trajectory.append(state_pred[0:3].copy())
                
                return np.array([thrust, roll_rate, pitch_rate, yaw_rate]), np.array(predicted_trajectory)
        
        # 返回安全控制输入 (保持悬停)
        state_pred = current_state.copy()
        u_hover = np.array([mass * g, 0, 0, 0])
        for i in range(N):
            state_pred = discrete_dynamics(state_pred, u_hover, dt)
            predicted_trajectory.append(state_pred[0:3].copy())
        
        return u_hover, np.array(predicted_trajectory)

# 初始化仿真
current_state = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])  # 初始状态
states = [current_state.copy()]
controls = []
predicted_trajectories = []  # 存储预测轨迹

# 创建图形
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')
ax.set_xlim(-1, 6)
ax.set_ylim(-1, 6)
ax.set_zlim(0, 4)
ax.set_xlabel('X (m)')
ax.set_ylabel('Y (m)')
ax.set_zlabel('Z (m)')
ax.set_title('Drone NMPC Control with Obstacle Avoidance and Prediction')

# 绘制目标点
ax.scatter(target_pos[0], target_pos[1], target_pos[2], c='g', marker='*', s=200, label='Target')

# 绘制障碍物
for obs in obstacles:
    u = np.linspace(0, 2 * np.pi, 12)
    v = np.linspace(0, np.pi, 8)
    x = obs["radius"] * np.outer(np.cos(u), np.sin(v)) + obs["pos"][0]
    y = obs["radius"] * np.outer(np.sin(u), np.sin(v)) + obs["pos"][1]
    z = obs["radius"] * np.outer(np.ones(np.size(u)), np.cos(v)) + obs["pos"][2]
    ax.plot_surface(x, y, z, color='red', alpha=0.3)

# 添加障碍物图例
from matplotlib.patches import Patch
obstacle_legend = Patch(color='red', alpha=0.3, label='Obstacle')

# 初始化轨迹线
trajectory_line, = ax.plot([], [], [], 'b-', label='Actual Trajectory')
prediction_line, = ax.plot([], [], [], 'y--', label='Predicted Trajectory')
prediction_points = ax.scatter([], [], [], c='orange', s=20, alpha=0.7, label='Prediction Points')

# 添加图例
ax.legend(handles=[trajectory_line, prediction_line, prediction_points,
                  ax.scatter([], [], [], c='g', marker='*', s=200, label='Target'),
                  obstacle_legend])

# 添加文本显示
time_text = ax.text2D(0.05, 0.95, '', transform=ax.transAxes)
state_text = ax.text2D(0.05, 0.90, '', transform=ax.transAxes)
control_text = ax.text2D(0.05, 0.85, '', transform=ax.transAxes)
prediction_text = ax.text2D(0.05, 0.80, '', transform=ax.transAxes)

# 初始化无人机位置标记
drone_marker = ax.scatter([], [], [], c='blue', s=100)

# 简化仿真循环
max_frames = 100  # 限制仿真帧数
frame_count = 0

def update(frame):
    global current_state, states, controls, predicted_trajectories, frame_count
    
    frame_count += 1
    if frame_count > max_frames:
        logger.info("Reached maximum frame count")
        return trajectory_line, prediction_line, prediction_points, drone_marker, time_text, state_text, control_text, prediction_text
    
    # 记录开始时间
    start_time = time.time()
    
    # 调用NMPC控制器
    u_opt, predicted_traj = nmpc_controller(current_state, target_pos, obstacles)
    
    # 记录计算时间
    comp_time = time.time() - start_time
    logger.info("Computation time: %.3f s, Control: %s", comp_time, u_opt)
    
    # 保存预测轨迹
    predicted_trajectories.append(predicted_traj)
    
    # 应用控制输入
    current_state = discrete_dynamics(current_state, u_opt, dt)
    
    # 保存状态和控制输入
    states.append(current_state.copy())
    controls.append(u_opt.copy())
    
    # 更新轨迹
    states_array = np.array(states)
    trajectory_line.set_data(states_array[:, 0], states_array[:, 1])
    trajectory_line.set_3d_properties(states_array[:, 2])
    
    # 更新预测轨迹
    if len(predicted_traj) > 0:
        prediction_line.set_data(predicted_traj[:, 0], predicted_traj[:, 1])
        prediction_line.set_3d_properties(predicted_traj[:, 2])
        
        # 更新预测点
        prediction_points._offsets3d = (predicted_traj[:, 0], predicted_traj[:, 1], predicted_traj[:, 2])
    
    # 更新无人机位置
    drone_marker._offsets3d = ([current_state[0]], [current_state[1]], [current_state[2]])
    
    # 更新文本信息
    time_text.set_text('Time: %.1f s, Comp: %.3f s' % (frame * dt, comp_time))
    state_text.set_text('Position: (%.2f, %.2f, %.2f) m\nVelocity: (%.2f, %.2f, %.2f) m/s' % 
                       (current_state[0], current_state[1], current_state[2],
                        current_state[3], current_state[4], current_state[5]))
    control_text.set_text('Thrust: %.2f N\nRates: (%.2f, %.2f, %.2f) rad/s' % 
                         (u_opt[0], u_opt[1], u_opt[2], u_opt[3]))
    
    # 显示预测信息
    if len(predicted_traj) > 0:
        prediction_text.set_text('Prediction Horizon: %d steps\nFinal Prediction: (%.2f, %.2f, %.2f)' % 
                               (N, predicted_traj[-1, 0], predicted_traj[-1, 1], predicted_traj[-1, 2]))
    else:
        prediction_text.set_text('Prediction Horizon: %d steps\nNo prediction available' % N)
    
    # 检查是否到达目标
    pos_error = np.linalg.norm(current_state[0:3] - target_pos)
    if pos_error < 0.5:
        logger.info("Reached target!")
        return trajectory_line, prediction_line, prediction_points, drone_marker, time_text, state_text, control_text, prediction_text
    
    return trajectory_line, prediction_line, prediction_points, drone_marker, time_text, state_text, control_text, prediction_text

# 创建动画
ani = FuncAnimation(fig, update, frames=max_frames, interval=dt*1000, blit=True, repeat=False)

plt.tight_layout()
plt.show()

# 仿真结束后打印总结
logger.info("Simulation finished")
logger.info("Final position: %s", current_state[0:3])
logger.info("Target position: %s", target_pos)
logger.info("Position error: %.3f m", np.linalg.norm(current_state[0:3] - target_pos))

# 绘制轨迹和障碍物的静态图
fig2 = plt.figure(figsize=(10, 8))
ax2 = fig2.add_subplot(111, projection='3d')

# 绘制轨迹
states_array = np.array(states)
ax2.plot(states_array[:, 0], states_array[:, 1], states_array[:, 2], 'b-', label='Actual Trajectory')

# 绘制预测轨迹 (最后几次预测)
for i, pred_traj in enumerate(predicted_trajectories[-5:]):  # 只显示最后5次预测
    alpha = 0.2 + 0.8 * (i / min(5, len(predicted_trajectories)))
    ax2.plot(pred_traj[:, 0], pred_traj[:, 1], pred_traj[:, 2], 'y--', alpha=alpha, label='Predicted Trajectory' if i == 0 else "")

# 绘制目标点
ax2.scatter(target_pos[0], target_pos[1], target_pos[2], c='g', marker='*', s=200, label='Target')

# 绘制障碍物
for obs in obstacles:
    u = np.linspace(0, 2 * np.pi, 12)
    v = np.linspace(0, np.pi, 8)
    x = obs["radius"] * np.outer(np.cos(u), np.sin(v)) + obs["pos"][0]
    y = obs["radius"] * np.outer(np.sin(u), np.sin(v)) + obs["pos"][1]
    z = obs["radius"] * np.outer(np.ones(np.size(u)), np.cos(v)) + obs["pos"][2]
    ax2.plot_surface(x, y, z, color='red', alpha=0.3)

ax2.set_xlim(-1, 6)
ax2.set_ylim(-1, 6)
ax2.set_zlim(0, 4)
ax2.set_xlabel('X (m)')
ax2.set_ylabel('Y (m)')
ax2.set_zlabel('Z (m)')
ax2.set_title('Drone Flight Path with Obstacles and Predictions')
ax2.legend()

plt.tight_layout()
plt.show()

# 绘制控制输入随时间的变化
fig3, axs = plt.subplots(2, 2, figsize=(12, 8))
controls_array = np.array(controls)
time_array = np.arange(len(controls)) * dt

axs[0, 0].plot(time_array, controls_array[:, 0])
axs[0, 0].set_title('Thrust')
axs[0, 0].set_ylabel('N')
axs[0, 0].grid(True)

axs[0, 1].plot(time_array, controls_array[:, 1])
axs[0, 1].set_title('Roll Rate')
axs[0, 1].set_ylabel('rad/s')
axs[0, 1].grid(True)

axs[1, 0].plot(time_array, controls_array[:, 2])
axs[1, 0].set_title('Pitch Rate')
axs[1, 0].set_ylabel('rad/s')
axs[1, 0].grid(True)

axs[1, 1].plot(time_array, controls_array[:, 3])
axs[1, 1].set_title('Yaw Rate')
axs[1, 1].set_ylabel('rad/s')
axs[1, 1].set_xlabel('Time (s)')
axs[1, 1].grid(True)

plt.tight_layout()
plt.show()
